from app import Factor, ResultAnalyse
import pandas as pd
import sys, os
from pyecharts.charts import Page

def calculate_factor(df, code):
    # 预期收益
    df["expect_close"] = df["close"].shift(-5)
    df["return"] = (df["expect_close"] - df["close"]).fillna(0)
    # 计算因子
    short_window = 5
    long_window = 10
    df["close_ema_short"] = df["close"].ewm(span=short_window, adjust=False).mean()
    df["close_ema_long"] = df["close"].ewm(span=long_window, adjust=False).mean()
    df["avg"] = df["amount"]/df["volume"]
    df["highxlow"] = df["low"] * df["high"]
    df['highxlow2'] = df['highxlow'].apply(lambda x: x**0.5)
    df['factor'] = df["highxlow2"] - df["avg"]
    # 计算相关性
    factor = Factor.Factor()
    factor.add_expect_column(df["return"])
    factor.add_factor_column(df["factor"])
    corr_value = factor.calculate_corr()
    print(corr_value)

    date_list = df["dateint"].tolist()
    time_list = [str(i) for i in date_list]
    expect_list = df["return"].tolist()
    factor_list = df["factor"].tolist()
    result_analyse = ResultAnalyse.ResultAnalyse()
    page = Page()
    # 净值
    yaxis = {
        "期望":expect_list, 
        "因子":factor_list, 
    }
    money_line = result_analyse.draw_line_with_multiple(time_list, yaxis, code, "300px")
    page.add(
        money_line
    )
    page.render("factor_{}.html".format(code))

def calculate_by_code():
    args = sys.argv[1:]
    folder = "data/data"
    file_name = args[0]

    csv_path = f"{folder}/{file_name}.csv"
    df = pd.read_csv(csv_path)
    calculate_factor(df, file_name)

def calculate_by_folder():
    args = sys.argv[1:]
    folder = args[0]
    file_names = os.listdir(folder)
    
    # 打印文件名
    for file_name in file_names:
        csv_path = f"{folder}/{file_name}"
        df = pd.read_csv(csv_path)
        calculate_factor(df)
    
if __name__ == '__main__':
    # calculate_by_folder()
    calculate_by_code()